import tank
import shotgun_api3
import pymel.core as pm
import yaml

def load_yaml_data(file_path):
    with open(file_path, 'r') as f:
        data = yaml.load(f)
    return data

def get_asset_instance():
    """
    # 1. 获取资产的shaderNetwork和所有abc的ref节点
    # 2. 获取资产对应shaderNetwork的shaderMapping数据
    return:
    {u'shaderMappingTest': {'id': 2201,
                            'items': [u'shaderMappingTest_Model_v001RN',
                                      u'shaderMappingTest_Model_v001RN1'],
                            'shader_network': u'shaderMappingTest_Shading_v001RN',
                            'shader_mapping': path},
    }
    """
    asset_dict = {}
    all_data = get_items()
    for data in all_data:
        asset_name = data["fields"]["Asset"]
        if asset_name not in asset_dict:
            asset_id = data["sg_data"]["entity"]["id"]
            asset_dict[asset_name] = {'id':asset_id}
            asset_dict[asset_name]['items'] = []
            asset_dict[asset_name]["shader_network"] = []
            asset_dict[asset_name]["shader_mapping"] = []
        asset_ref_node = data["node_name"]
        asset_publish = data["sg_data"]["published_file_type"]["name"]
        if asset_publish == "Shader Network":
            # 如果有多个shader文件，只使用第一个 （To Do ）
            if asset_dict[asset_name]["shader_network"]:
                continue
            asset_dict[asset_name]["shader_network"] = asset_ref_node

            # 获取最高版本的关系文件
            asset_shader_mapping_data = get_asset_published_files(asset_id=asset_dict[asset_name]["id"], publish_type="Shader Mapping")
            asset_shader_mapping = sorted(asset_shader_mapping_data, key=lambda version: version["version_number"])[-1]
            asset_dict[asset_name]["shader_mapping"] = asset_shader_mapping["path"]['local_path_windows']
        else:
            asset_dict[asset_name]["items"].append(asset_ref_node)
    return asset_dict
        
def assign_material():
    asset_data = get_asset_instance()
    
    for asset in asset_data:
        namespace_list = []
        engine_list = []
        # 获取资产的所有namespace
        data = asset_data[asset]
        for ref in data["items"]:
            namespace_list.append(pm.PyNode(ref).referenceFile().namespace)
            
        # 获取shaderNetwork的namespace
        shader_network_ref = data["shader_network"]
        shader_namespace = pm.PyNode(shader_network_ref).referenceFile().namespace
        
        # 获取关系数据 
        shader_mapping_file = data["shader_mapping"]
        shader_mapping_data = load_yaml_data(shader_mapping_file)
        
        # 解析namespace并赋予材质
        for engine in shader_mapping_data:
            final_engine = ':'.join([shader_namespace,engine])
            print final_engine
            element_list = shader_mapping_data[engine]
            final_element_list = [element.replace("|", "|{}:".format(namespace)) for namespace in namespace_list for element in element_list]
            pm.sets(final_engine, forceElement=final_element_list)


def mapping_gemotry(engine, element):
    pm.sets(engine, forceElement=element)
    
    
def get_asset_published_files(asset_id, publish_type):
    """
    get_asset_published_files(asset_id=2201, publish_type="Shader Mapping")
    args:
        asset_id: 
        publish_type: 
    retrun: list.  e.g.
    [
    {'id': 6445,
    'path': {'content_type': None,
           'id': 62784,
           'link_type': 'local',
           'local_path': 'X:\\thingdom\\assets\\character\\shaderMappingTest\\shaders\\shaderMappingTest.Shading.v004.yml',
           'local_path_linux': '/nas/Library/3dassets/thingdom/assets/character/shaderMappingTest/shaders/shaderMappingTest.Shading.v004.yml',
           'local_path_mac': '/nas/Library/3dassets/thingdom/assets/character/shaderMappingTest/shaders/shaderMappingTest.Shading.v004.yml',
           'local_path_windows': 'X:\\thingdom\\assets\\character\\shaderMappingTest\\shaders\\shaderMappingTest.Shading.v004.yml',
           'local_storage': {'id': 3,
                             'name': '3dassetRoot',
                             'type': 'LocalStorage'},
           'name': 'shaderMappingTest.Shading.v004.yml',
           'type': 'Attachment',
           'url': 'file://X:\\thingdom\\assets\\character\\shaderMappingTest\\shaders\\shaderMappingTest.Shading.v004.yml'},
    'type': 'PublishedFile',
    'version_number': 4},

    ......
    
    ] 
    
    """

    sg = shotgun_api3.Shotgun("http://shotgun-staging.3dbox.com",
                              login='shenmeng',
                              password='sheen@3dbox.com')
    filters = [
        ["entity", "is", {"id":asset_id, 'type': 'Asset'}],
        #["name", "is", "shaderMappingTest"],
        ["published_file_type", "name_is", publish_type], 
    ]

    fields = ["version_number", "path"]
    data = sg.find('PublishedFile', filters, fields)
    return data
    
def get_items():
    items = []
    g_cached_sg_publish_data = {}
    engine = tank.platform.current_engine()
    scene_objects = scan_scene()
    for scene_object in scene_objects:
        node_name = scene_object.get("node")
        node_type = scene_object.get("type")
        file_name = scene_object.get("path").replace("/", os.path.sep)
        
        matching_template = engine.tank.template_from_path(file_name)
        if matching_template:
            fields = matching_template.get_fields(file_name)
            if 'version' in fields:
                for key_name, key in matching_template.keys.iteritems():
                    if key_name in fields and key.is_abstract:
                        del(fields[key_name])
                # we also want to normalize the eye field (this should probably be an abstract field!)
                # note: we need to do this explicitly because the eye isn't abstract in the default 
                # configs yet (which is incorrect!).
                fields["eye"] = "%V"
                # now build the normalized path that we can use to find corresponding Shotgun published files
                normalized_path = matching_template.apply_fields(fields)

                item = {}
                item["path"] = normalized_path
                item["node_name"] = node_name
                item["node_type"] = node_type
                #item["template"] = matching_template
                item["fields"] = fields
                item["sg_data"] = None

                # store the normalized fields in dict
                items.append(item)
                
    valid_paths = [ x.get("path") for x in items ]
    # check if we have the path in the cache
    paths_to_fetch = []
    for p in valid_paths:
        if p not in g_cached_sg_publish_data:
            paths_to_fetch.append(p)
        else:
            # use cache data!
            for item in items:
                if item.get("path") == p:
                    item["sg_data"] = g_cached_sg_publish_data[p]
    fields = ["entity",
              "entity.Asset.sg_asset_type", # grab asset type if it is an asset
              #"code",
              #"image",
              #"name",
              #"task",
              #"version_number",
              #"project"
              ]
    if tank.util.get_published_file_entity_type(engine.tank) == "PublishedFile":
        fields.append("published_file_type")
    else:# == "TankPublishedFile"
        fields.append("tank_type")
    sg_data = tank.util.find_publish(engine.tank, paths_to_fetch, fields=fields)
    # process and cache shotgun items
    for (path, sg_chunk) in sg_data.items():
        # cache item
        g_cached_sg_publish_data[path] = sg_chunk
        # append the sg data to the right path
        for item in items:
            if item.get("path") == path:
                item["sg_data"] = sg_chunk
        
    return items
    
def scan_scene():
    refs = []
    # first let's look at maya references
    for x in pm.listReferences():
        node_name = x.refNode.longName()
        # get the path and make it platform dependent
        # (maya uses C:/style/paths)
        maya_path = x.path.replace("/", os.path.sep)
        refs.append( {"node": node_name, "type": "reference", "path": maya_path})
    return refs